# video_effects.py
from moviepy.editor import *
import numpy as np
from PIL import Image, ImageEnhance
import cv2
import os

class VideoEffects:
    def __init__(self):
        self.effects = {
            "brightness": self.adjust_brightness,
            "contrast": self.adjust_contrast,
            "saturation": self.adjust_saturation,
            "blur": self.apply_blur,
            "sharpen": self.apply_sharpen,
            "edge_detection": self.apply_edge_detection
        }
    
    def add_particle_effect(self, video_clip, effect_type: str = "snow", density: int = 50):
        """添加粒子效果"""
        def make_frame(t):
            # 获取原帧
            frame = video_clip.get_frame(t)
            h, w = frame.shape[:2]
            
            # 创建一个与原帧相同大小的覆盖层
            overlay = frame.copy()
            
            # 生成随机粒子位置
            for _ in range(density):
                x = np.random.randint(0, w)
                y = np.random.randint(0, h)
                size = np.random.randint(1, 3)
                
                if effect_type == "snow":
                    color = (255, 255, 25)  # 白色雪花
                elif effect_type == "rain":
                    color = (100, 100, 200)  # 深蓝色雨滴
                elif effect_type == "stars":
                    color = (255, 255, 200)  # 星星
                else:
                    color = (255, 255, 255)  # 默认白色
                
                # 在帧上绘制粒子
                cv2.circle(overlay, (x, y), size, color, -1)
            
            # 混合原帧和粒子层
            alpha = 0.3  # 粒子透明度
            blended = cv2.addWeighted(frame, 1 - alpha, overlay, alpha, 0)
            
            return blended
        
        particle_clip = VideoClip(make_frame, duration=video_clip.duration)
        return CompositeVideoClip([video_clip, particle_clip])
    
    def apply_filter(self, video_clip, filter_type: str):
        """应用滤镜"""
        filters = {
            "vintage": self._vintage_filter,
            "cinematic": self._cinematic_filter,
            "warm": self._warm_filter,
            "cool": self._cool_filter,
            "black_and_white": self._black_and_white_filter,
            "sepia": self._sepia_filter
        }
        filter_func = filters.get(filter_type, lambda x: x)
        return filter_func(video_clip)
    
    def _vintage_filter(self, video_clip):
        """复古滤镜"""
        def filter_frame(image):
            img = Image.fromarray(image)
            
            # 调整对比度和饱和度
            enhancer = ImageEnhance.Contrast(img)
            img = enhancer.enhance(1.2)
            
            enhancer = ImageEnhance.Color(img)
            img = enhancer.enhance(0.8)
            
            # 添加暖色调
            img_array = np.array(img)
            img_array = img_array.astype(np.float64)
            
            # 增强红色通道，减弱蓝色通道
            img_array[:, :, 0] = np.clip(img_array[:, :, 0] * 1.1, 0, 255)  # 红色通道增强
            img_array[:, :, 2] = np.clip(img_array[:, :, 2] * 0.9, 0, 255)  # 蓝色通道减弱
            
            return img_array.astype(np.uint8)
        
        return video_clip.fl_image(filter_frame)
    
    def _cinematic_filter(self, video_clip):
        """电影滤镜"""
        def filter_frame(image):
            img = image.astype(np.float64)
            
            # 调整对比度
            img = 128 + 1.2 * (img - 128)
            
            # 调整色彩平衡
            img[:, :, 0] *= 0.95  # 红色
            img[:, :, 1] *= 1.0   # 绿色
            img[:, :, 2] *= 1.1   # 蓝色
            
            # 限制值范围
            img = np.clip(img, 0, 255)
            
            return img.astype(np.uint8)
        
        return video_clip.fl_image(filter_frame)
    
    def _warm_filter(self, video_clip):
        """暖色调滤镜"""
        def filter_frame(image):
            img = image.astype(np.float64)
            
            # 增强红色和黄色
            img[:, :, 0] = np.clip(img[:, :, 0] * 1.1, 0, 255)  # 红色
            img[:, :, 1] = np.clip(img[:, :, 1] * 1.05, 0, 255)  # 绿色
            img[:, :, 2] = np.clip(img[:, :, 2] * 0.9, 0, 255)   # 蓝色
            
            return img.astype(np.uint8)
        
        return video_clip.fl_image(filter_frame)
    
    def _cool_filter(self, video_clip):
        """冷色调滤镜"""
        def filter_frame(image):
            img = image.astype(np.float64)
            
            # 增强蓝色
            img[:, :, 0] = np.clip(img[:, :, 0] * 0.9, 0, 255)   # 红色
            img[:, :, 1] = np.clip(img[:, :, 1] * 0.95, 0, 255)  # 绿色
            img[:, :, 2] = np.clip(img[:, :, 2] * 1.1, 0, 255)   # 蓝色
            
            return img.astype(np.uint8)
        
        return video_clip.fl_image(filter_frame)
    
    def _black_and_white_filter(self, video_clip):
        """黑白滤镜"""
        def filter_frame(image):
            # 转换为灰度
            gray = np.dot(image[...,:3], [0.299, 0.587, 0.114])
            gray = gray.astype(np.uint8)
            # 扩展为3通道
            return np.stack([gray, gray, gray], axis=-1)
        
        return video_clip.fl_image(filter_frame)
    
    def _sepia_filter(self, video_clip):
        """棕褐色滤镜"""
        def filter_frame(image):
            img = image.astype(np.float64)
            
            # 应用棕褐色变换矩阵
            sepia_filter = np.array([
                [0.393, 0.769, 0.189],
                [0.349, 0.686, 0.168],
                [0.272, 0.534, 0.131]
            ])
            
            # 应用变换
            sepia_img = img @ sepia_filter.T
            sepia_img = np.clip(sepia_img, 0, 255)
            
            return sepia_img.astype(np.uint8)
        
        return video_clip.fl_image(filter_frame)
    
    def adjust_brightness(self, video_clip, factor: float = 1.0):
        """调整亮度 (factor > 1.0 增加亮度, factor < 1.0 降低亮度)"""
        def adjust_frame(image):
            img = image.astype(np.float64)
            img = img * factor
            img = np.clip(img, 0, 255)
            return img.astype(np.uint8)
        
        return video_clip.fl_image(adjust_frame)
    
    def adjust_contrast(self, video_clip, factor: float = 1.0):
        """调整对比度"""
        def adjust_frame(image):
            img = image.astype(np.float64)
            # 中心值为128
            img = 128 + factor * (img - 128)
            img = np.clip(img, 0, 255)
            return img.astype(np.uint8)
        
        return video_clip.fl_image(adjust_frame)
    
    def adjust_saturation(self, video_clip, factor: float = 1.0):
        """调整饱和度"""
        def adjust_frame(image):
            img = Image.fromarray(image)
            enhancer = ImageEnhance.Color(img)
            img = enhancer.enhance(factor)
            return np.array(img)
        
        return video_clip.fl_image(adjust_frame)
    
    def apply_blur(self, video_clip, radius: float = 1.0):
        """应用模糊效果"""
        from PIL import ImageFilter
        def blur_frame(image):
            img = Image.fromarray(image)
            img = img.filter(ImageFilter.GaussianBlur(radius=radius))
            return np.array(img)
        
        return video_clip.fl_image(blur_frame)
    
    def apply_sharpen(self, video_clip):
        """应用锐化效果"""
        from PIL import ImageFilter
        def sharpen_frame(image):
            img = Image.fromarray(image)
            img = img.filter(ImageFilter.SHARPEN)
            return np.array(img)
        
        return video_clip.fl_image(sharpen_frame)
    
    def apply_edge_detection(self, video_clip):
        """应用边缘检测"""
        import cv2
        def edge_frame(image):
            gray = cv2.cvtColor(image, cv2.COLOR_RGB2GRAY)
            edges = cv2.Canny(gray, 100, 200)
            # 将边缘转换为三通道
            edges_3ch = cv2.cvtColor(edges, cv2.COLOR_GRAY2RGB)
            return edges_3ch
        
        return video_clip.fl_image(edge_frame)
    
    def apply_effect(self, video_clip, effect_type: str, **kwargs):
        """应用指定效果"""
        effect_func = self.effects.get(effect_type)
        if effect_func:
            return effect_func(video_clip, **kwargs)
        else:
            return video_clip
    
    def add_motion_blur(self, video_clip, angle: float = 0, length: int = 10):
        """添加运动模糊效果"""
        def motion_blur_frame(image):
            # 创建运动模糊核
            kernel = np.zeros((length, length))
            kernel[int((length-1)/2), :] = np.ones(length)
            kernel = kernel / length
            
            # 旋转核以匹配角度
            if angle != 0:
                # 简化处理，实际应用中可能需要更复杂的旋转
                pass
            
            # 应用卷积
            img = image.astype(np.float64)
            for i in range(3):  # 对每个颜色通道应用模糊
                img[:, :, i] = cv2.filter2D(img[:, :, i], -1, kernel)
            
            return img.astype(np.uint8)
        
        return video_clip.fl_image(motion_blur_frame)
    
    def add_fade_in(self, video_clip, duration: float = 1.0):
        """添加淡入效果"""
        return video_clip.fadein(duration)
    
    def add_fade_out(self, video_clip, duration: float = 1.0):
        """添加淡出效果"""
        return video_clip.fadeout(duration)
